package io.gitee.tianhaoran.chat.sdk.provider.chatgpt;

import cn.hutool.http.ContentType;
import cn.hutool.http.Header;
import com.alibaba.fastjson2.JSON;
import com.dtflys.forest.Forest;
import com.dtflys.forest.http.ForestRequest;
import io.gitee.tianhaoran.chat.sdk.provider.Convertor;
import io.gitee.tianhaoran.chat.sdk.provider.chatgpt.config.ChatgptConfig;
import io.gitee.tianhaoran.chat.sdk.provider.chatgpt.model.ChatgptImageRequest;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import io.gitee.tianhaoran.chat.sdk.api.AbstractChatBlender;
import io.gitee.tianhaoran.chat.sdk.api.Response;
import io.gitee.tianhaoran.chat.sdk.api.ValidProvider;
import io.gitee.tianhaoran.chat.sdk.api.entity.chat.CompletionRequest;
import io.gitee.tianhaoran.chat.sdk.api.entity.image.ImageRequest;
import io.gitee.tianhaoran.chat.sdk.api.group.ChatgptGroup;
import io.gitee.tianhaoran.chat.sdk.provider.chatgpt.model.ChatgptCompletionRequest;
import io.gitee.tianhaoran.chat.sdk.provider.chatgpt.model.ChatgptCompletionResponse;
import io.gitee.tianhaoran.chat.sdk.provider.chatgpt.model.ChatgptImageResponse;

import java.util.concurrent.atomic.AtomicReference;

import static io.gitee.tianhaoran.chat.sdk.api.def.Const.DEFAULT_USER_AGENT;
import static io.gitee.tianhaoran.chat.sdk.provider.chatgpt.config.Const.Chatgpt_Completion_Uri;
import static io.gitee.tianhaoran.chat.sdk.provider.chatgpt.config.Const.Chatgpt_Image_Uri;

/**
 * Chatgpt 厂商实现
 * <p>
 * 作者：沈自在
 **/
@Slf4j
public class ChatgptBlender extends AbstractChatBlender {

    /**
     * chatgpt配置
     */
    private final ChatgptConfig config;

    /**
     *
     */
    private final EventSource.Factory factory;

    /**
     * Api 地址
     */
    @Getter
    public final String Api_Host;


    public ChatgptBlender(ChatgptConfig config) {
        this.config = config;
        this.factory = createRequestFactory();
        Api_Host = (config.getIsProxy() ? config.getProxyUrl() : config.getHost());
    }


    @Override
    public Response completions(CompletionRequest request) {
        return completions(Api_Host, config.getKey(), request);
    }

    @Override
    public Response completions(String host, String key, CompletionRequest request) {
        // 1. 预处理
        boolean isValid = ValidProvider.validate(request, ChatgptGroup.class);
        AtomicReference<ChatgptCompletionResponse> gptResponse = new AtomicReference<>();
        if (isValid && !request.isStream()) {
            // 2. 转换请求
            ChatgptCompletionRequest gptRequest = Convertor.toGptRequestWithChat(request);

            // 3. 初始化 Forest 请求
            ForestRequest<?> forestRequest = Forest.post(Chatgpt_Completion_Uri)
                    .basePath(host)
                    .contentTypeJson()
                    .addHeader(Header.AUTHORIZATION.getValue(), "Bearer " + key)
                    .addHeader(Header.USER_AGENT.getValue(), DEFAULT_USER_AGENT)
                    .onSuccess((data, req, res) -> {
                        gptResponse.set(JSON.parseObject(res.getContent(), ChatgptCompletionResponse.class));
                    })
                    .onError((ex, req, res) -> {
                        log.error("[AI-Chat] error sending chatgpt request, occur {}", ex.getMessage());
                    })
                    .addBody(JSON.toJSONString(gptRequest));

            // 6. 请求执行
            try {
                forestRequest.execute();
            } catch (Exception ex) {
                log.error("[AI-Chat] error sending chatgpt request, occur {}", ex.getMessage());
            }
        }
        return gptResponse.get();
    }

    @Override
    public EventSource completions(CompletionRequest request, EventSourceListener listener) {
        return completions(Api_Host, config.getKey(), request, listener);
    }

    @Override
    public EventSource completions(String host, String key, CompletionRequest completionRequest, EventSourceListener listener) {
        boolean isValid = ValidProvider.validate(completionRequest, ChatgptGroup.class);

        completionRequest.setStream(Boolean.TRUE);
        //if (!completionRequest.isStream()) {
        //    throw new ChatException("You are calling the stream question answering interface, but the stream argument is indeed false");
        //}

        // 校验通过
        if (isValid) {
            // 1. 转换请求
            ChatgptCompletionRequest gptRequest = Convertor.toGptRequestWithChat(completionRequest);

            // 2. 获取地址
            Request request = new Request.Builder()
                    .header(Header.AUTHORIZATION.getValue(), "Bearer " + key)
                    .url(host.concat(Chatgpt_Completion_Uri))
                    .post(RequestBody.create(MediaType.parse(ContentType.JSON.getValue()), JSON.toJSONString(gptRequest)))
                    .build();

            // 3.
            return factory.newEventSource(request, listener);
        }
        return null;
    }

    @Override
    public Response generateImage(ImageRequest request) {
        return generateImage(Api_Host, config.getKey(), request);
    }

    @Override
    public Response generateImage(String host, String key, ImageRequest request) {
        boolean isValid = ValidProvider.validate(request, ChatgptGroup.class);
        AtomicReference<ChatgptImageResponse> gptResponse = new AtomicReference<>();
        if (isValid) {
            ChatgptImageRequest gptRequest = Convertor.toGptRequestWithImage(request);
            ForestRequest<?> forestRequest = Forest.post(Chatgpt_Image_Uri)
                    .basePath(host)
                    .contentTypeJson()
                    .addHeader(Header.AUTHORIZATION.getValue(), "Bearer " + key)
                    .addHeader(Header.USER_AGENT.getValue(), DEFAULT_USER_AGENT)
                    .onSuccess((data, req, res) -> {
                        gptResponse.set(JSON.parseObject(res.getContent(), ChatgptImageResponse.class));
                    })
                    .onError((ex, req, res) -> {
                        log.error("[AI-Chat] error sending chatgpt request, occur {}", ex.getMessage());
                    })
                    .addBody(JSON.toJSONString(gptRequest));

            try {
                forestRequest.execute();
            } catch (Exception ex) {
                log.error("[AI-Chat] error sending chatgpt request, occur {}", ex.getMessage());
            }
        }
        return gptResponse.get();
    }
}
